import tensorflow as tf
import sys
import numpy as np

tf.random.set_seed(777)

iters = 300
alpha = 0.01
group = 20


def sep(label):
    print('-' * 32, label, '-' * 32, sep='')


# (2)	求f(x) = a*x**2 + b*x + c的最小值
# ①	合理创建变量和常量
x = tf.Variable(0, dtype=tf.float32, name='x')
a = tf.constant(2, dtype=tf.float32, name='a')
b = tf.constant(3, dtype=tf.float32, name='b')
c = tf.constant(4, dtype=tf.float32, name='c')


@tf.function
def get_y():
    y = a * x ** 2 + b * x + c
    return y


opt = tf.keras.optimizers.Adam(name='opt', learning_rate=alpha)


# ②	设定优化模型
@tf.function
def xstep():
    opt.minimize(get_y, [x])


# ③	循环迭代计算
last_x = x.numpy()
for i in range(iters):
    xstep()
    this_x = x.numpy()
    if i % group == 0:
        print(f'#{i + 1}: x = {this_x}')
    if np.isclose(this_x, last_x):
        print('Converged!')
        break
    last_x = this_x
if i % group != 0:
    print(f'#{i + 1}: x = {this_x}')

# ④	打印相关数据值
print(f'理论最小值: {(- b / (2*a)).numpy()}')
print(f'模型迭代得到最小值: {this_x}')
